package cn.donting.jme.utile;


import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.util.HashMap;
import java.util.Map;

public class IdentityWeakHashMap<K, V> {
    private HashMap<HashWeakReference<K>, V> identityHashMap = new HashMap<>();
    private final ReferenceQueue<K> referenceQueue = new ReferenceQueue<>();

    public V put(K k, V v) {
        cleanup();
        HashWeakReference<K> weakKey = new HashWeakReference<>(k, referenceQueue);
        return identityHashMap.put(weakKey, v);
    }

    public V get(K key) {
        cleanup();
        HashWeakReference<K> weakKey = new HashWeakReference<>(key);
        return identityHashMap.get(weakKey);
    }

    public V remove(K key) {
        cleanup();
        HashWeakReference<K> weakKey = new HashWeakReference<>(key);
        return identityHashMap.remove(weakKey);
    }

    private void cleanup() {
        WeakReference<K> ref;
        while ((ref = (WeakReference<K>) referenceQueue.poll()) != null) {
            identityHashMap.remove(ref);
        }
    }

    private static class HashWeakReference<T> extends WeakReference<T> {
        private int hashCode;

        public HashWeakReference(T referent) {
            super(referent);
            hashCode = referent.hashCode();
        }

        public HashWeakReference(T referent, ReferenceQueue<? super T> q) {
            super(referent, q);
            hashCode = referent.hashCode();
        }

        @Override
        public int hashCode() {
            return hashCode;
        }

        @Override
        public boolean equals(Object obj) {
            return ((HashWeakReference<?>) obj).get() == get();
        }
    }


    public static void main(String[] args) {
        IdentityWeakHashMap<String, String> stringIdentityWeakHashMap = new IdentityWeakHashMap<>();
        String s1 = new String("1");
        stringIdentityWeakHashMap.put(s1, "1");
        stringIdentityWeakHashMap.put(new String("1"), "2");

        System.gc();

        String s = stringIdentityWeakHashMap.get(new String("1"));
        String s2 = stringIdentityWeakHashMap.get(s1);

        System.out.println();
    }
}
